import cv2
import insightface
from insightface.app import FaceAnalysis 
from insightface.data import get_image as ins_get_image
import numpy as np

def draw_with_bbox():
    app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
    app.prepare(ctx_id=0, det_size=(320, 320))
    img = cv2.imread("tmp/person1.png")
    faces = app.get(img)
    rimg = app.draw_on(img, faces)
    cv2.imwrite("tmp/rec_person.png", rimg)

def swap_face():
    app = FaceAnalysis()
    app.prepare(ctx_id=0, det_size=(640, 640))    
    swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True)
    source_img = cv2.imread("tmp/person.jpeg")
    source_faces = app.get(source_img)
    source_face = source_faces[0]
    target_img = cv2.imread("tmp/person1.png")
    target_faces = app.get(target_img)
    target_face = target_faces[0]
    res = target_img.copy()
    res = swapper.get(res, target_face, source_face, paste_back=True)
    cv2.imwrite("tmp/swapface.png", res)
    '''
    res = []
    for face in target_faces:
        _img, _ = swapper.get(target_img, face, source_face, paste_back=False)
        res.append(_img)
    res = np.concatenate(res, axis=1)
    cv2.imwrite("tmp/swap.png", res)
    '''

if __name__ == '__main__':
    swap_face()
